import os
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import PIL.Image, PIL.ImageOps
import matplotlib.pyplot as plt
import math


class KeypointDataset(Dataset):
    def __init__(self,
                 image_dir,
                 label_dir,
                 transform=None,
                 target_input_size=(256, 256),
                 target_output_size=(256, 256),
                 sigma=4):
        super(KeypointDataset, self).__init__()
        self.image_dir = image_dir # 图片文件夹路径
        self.label_dir = label_dir # Label文件夹路径
        self.image_files = [f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.png'))] # 图片名
        self.transform = transform # 图片转换为Tensor矩阵
        self.target_input_size = target_input_size # 输入图片放缩的尺寸
        self.target_output_size = target_output_size # 输出heatmap的尺寸
        self.sigma = sigma # 高斯圈的sigma
        self.image_paths = [] # 每张图片的路径

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        image_path = os.path.join(self.image_dir, self.image_files[idx])
        self.image_paths.append(image_path)
        label_path = os.path.join(self.label_dir, self.image_files[idx].replace('.jpg', '.txt').replace('.png', '.txt'))
        radius = 0 # 初始化高斯圈的半径

        # 读取图像
        image = PIL.Image.open(image_path)
        # 如果图片格式为RGBA,则只保留前三个通道
        if image.mode == 'RGBA':
            image = image.convert("RGB")

        # 对输入图片进行缩放填充
        input_image, _, _ = resize_and_padding(image, self.target_input_size)
        # 将填充后的图片转换为灰度图（为什么在填充后转灰度是因为考虑到之后如果需要三通道方便修改）
        input_image = input_image.convert("L")
        # 将图片转为Tensor
        if self.transform:
            input_image = self.transform(input_image)

        # 获取输出heatmap的放缩填充比例，为计算ground truth的heatmap做准备
        _, scale, padding = resize_and_padding(image, self.target_output_size)

        # 读取对应的.txt文件中的标签
        with open(label_path, 'r') as file:
            lines = file.readlines()
            labels = [list(map(float, line.strip().split())) for line in lines]

        # 调整关键点坐标
        adjusted_labels = adjust_keypoints(labels, scale, padding)


        # code128系列（包括code39)的热图初始化
        # 初始化CodeX的两张主要热图：一张用于第1和第3个角点，另一张用于第2和第4个角点
        heatmap_codeX_13 = np.zeros((self.target_output_size[1], self.target_output_size[0]))
        heatmap_codeX_24 = np.zeros((self.target_output_size[1], self.target_output_size[0]))
        # 初始化方向热图：一张用于sin值，一张用于cos值
        heatmap_codeX_sin = np.full((self.target_output_size[1], self.target_output_size[0]), -2.0)
        heatmap_codeX_cos = np.full((self.target_output_size[1], self.target_output_size[0]), -2.0)
        # 初始化权重图，这里简化为两张，对应上面的热图
        weight_codeX_13 = np.ones((self.target_output_size[1], self.target_output_size[0]))
        weight_codeX_24 = np.ones((self.target_output_size[1], self.target_output_size[0]))
        weight_codeX_sin = np.ones((self.target_output_size[1], self.target_output_size[0]))  # 对应角度信息sin的权重矩阵
        weight_codeX_cos = np.ones((self.target_output_size[1], self.target_output_size[0]))  # 对应角度信息cos的权重矩阵

        # QR码的热图初始化
        # 初始化两张主要的热图：一张用于三个回字形角点，一张用于第三个角点
        heatmap_QR_124 = np.zeros((self.target_output_size[1], self.target_output_size[0]))
        heatmap_QR_3 = np.zeros((self.target_output_size[1], self.target_output_size[0]))
        # 初始化方向热图：一张用于sin值，一张用于cos值
        heatmap_QR_sin = np.full((self.target_output_size[1], self.target_output_size[0]), -2.0)
        heatmap_QR_cos = np.full((self.target_output_size[1], self.target_output_size[0]), -2.0)
        # 初始化权重图，这里简化为两张，对应上面的热图
        weight_QR_124 = np.ones((self.target_output_size[1], self.target_output_size[0]))
        weight_QR_3 = np.ones((self.target_output_size[1], self.target_output_size[0]))
        weight_QR_sin = np.ones((self.target_output_size[1], self.target_output_size[0]))  # 对应角度信息sin的权重矩阵
        weight_QR_cos = np.ones((self.target_output_size[1], self.target_output_size[0]))  # 对应角度信息cos的权重矩阵

        # DM码的热图初始化
        # 初始化两张主要的热图：一张用于有直边相连的角点，一张用于第二个角点
        heatmap_DM_134 = np.zeros((self.target_output_size[1], self.target_output_size[0]))
        heatmap_DM_2 = np.zeros((self.target_output_size[1], self.target_output_size[0]))
        # 初始化方向热图：一张用于sin值，一张用于cos值
        heatmap_DM_sin = np.full((self.target_output_size[1], self.target_output_size[0]), -2.0)
        heatmap_DM_cos = np.full((self.target_output_size[1], self.target_output_size[0]), -2.0)
        # 初始化权重图，这里简化为两张，对应上面的热图
        weight_DM_134 = np.ones((self.target_output_size[1], self.target_output_size[0]))
        weight_DM_2 = np.ones((self.target_output_size[1], self.target_output_size[0]))
        weight_DM_sin = np.ones((self.target_output_size[1], self.target_output_size[0]))  # 对应角度信息sin的权重矩阵
        weight_DM_cos = np.ones((self.target_output_size[1], self.target_output_size[0]))  # 对应角度信息cos的权重矩阵

        # ean的热图初始化
        # 初始化ean的两张主要热图：一张用于第1和第2个角点，另一张用于第3和第4个角点
        heatmap_EAN_12 = np.zeros((self.target_output_size[1], self.target_output_size[0]))
        heatmap_EAN_34 = np.zeros((self.target_output_size[1], self.target_output_size[0]))
        # 初始化方向热图：一张用于sin值，一张用于cos值
        heatmap_EAN_sin = np.full((self.target_output_size[1], self.target_output_size[0]), -2.0)
        heatmap_EAN_cos = np.full((self.target_output_size[1], self.target_output_size[0]), -2.0)
        # 初始化权重图，这里简化为两张，对应上面的热图
        weight_EAN_12 = np.ones((self.target_output_size[1], self.target_output_size[0]))
        weight_EAN_34 = np.ones((self.target_output_size[1], self.target_output_size[0]))
        weight_EAN_sin = np.ones((self.target_output_size[1], self.target_output_size[0]))  # 对应角度信息sin的权重矩阵
        weight_EAN_cos = np.ones((self.target_output_size[1], self.target_output_size[0]))  # 对应角度信息cos的权重矩阵


        def update_heatmaps_and_direction(x, y, target_coords,
                                          heatmap1, heatmap2, weight_map, weight,
                                          sin_map, cos_map, weight_sin, weight_cos,
                                          sigma):
            # 生成当前角点的高斯热图
            curr_heatmap, radius = generate_heatmap(np.zeros_like(heatmap1), (x, y), sigma)
            # 定义方向信息更新的热图阈值
            direction_threshold = 0.6  # 根据实际情况调整这个阈值
            # 遍历新高斯圈内所有点来更新方向信息
            indices = np.argwhere(curr_heatmap > 0)
            for index in indices:
                i, j = index[0], index[1]  # 获取新高斯圈内点的坐标
                # 只在以下两种情况下更新方向信息：
                # 1. 点在原热图中的值低于阈值（即没有被之前的重要高斯圈覆盖）。
                # 2. 点不在原热图的覆盖范围内。
                if heatmap1[i, j] <= direction_threshold and heatmap2[i, j] <= direction_threshold:
                    # 计算每个点到目标点的角度
                    angle_rad = calculate_angle([j, i], target_coords)
                    # 转换为正弦和余弦值
                    sin_val = np.sin(angle_rad)
                    cos_val = np.cos(angle_rad)
                    # 更新sin和cos值
                    sin_map[i, j] = sin_val
                    cos_map[i, j] = cos_val
            # 更新权重
            weight_sin[curr_heatmap > 0] = weight
            weight_cos[curr_heatmap > 0] = weight
            # 定义热图值的阈值
            threshold = 0.7
            # 更新主热图，特殊处理重叠区域
            overlap_indices = (curr_heatmap > 0) & (heatmap1 > 0)  # 定位重叠区域
            non_overlap_indices = (curr_heatmap > 0) & (heatmap1 == 0)  # 定位非重叠区域
            # 对于重叠区域，检查原热图或新热图中的值是否大于阈值
            high_value_indices = (heatmap1 > threshold) | (curr_heatmap > threshold)
            # 重叠且两者都不大于阈值的区域计算平均值
            average_indices = overlap_indices & (~high_value_indices)
            heatmap1[average_indices] = (heatmap1[average_indices] + curr_heatmap[average_indices]) / 2
            # 重叠且至少有一个大于阈值的区域，取较大值
            max_indices = overlap_indices & high_value_indices
            heatmap1[max_indices] = np.maximum(heatmap1[max_indices], curr_heatmap[max_indices])
            # 非重叠区域直接叠加
            heatmap1[non_overlap_indices] += curr_heatmap[non_overlap_indices]
            # 更新权重
            weight_map[curr_heatmap > 0] = weight

        # 遍历标签数据，每四行一组，处理每个QR码
        for i in range(0, len(adjusted_labels), 4):
            # 提取四个角点坐标
            corner_points = [adjusted_labels[i + k][:2] for k in range(4)]
            # 计算中心点坐标
            center_x, center_y = calculate_center_point(corner_points)
            # 使用中心点作为目标坐标
            target_coords = (center_x, center_y)

            first_point_coords = adjusted_labels[i][:5]  # 第一个角点
            # 如果是Code128系列
            if first_point_coords[2] == 0 or first_point_coords[2] == 3:
                # sigma = 2  # 高斯热图的标准差
                # 处理13角点
                for k in [0, 2]:  # 只遍历第1和第3个点
                    x, y, code_type, is_complete, weight = adjusted_labels[i + k][:5]
                    if is_complete:
                        update_heatmaps_and_direction(x, y, target_coords,
                                                      heatmap_codeX_13, heatmap_codeX_24, weight_codeX_13, weight,
                                                      heatmap_codeX_sin, heatmap_codeX_cos,
                                                      weight_codeX_sin, weight_codeX_cos, self.sigma)

                # 处理24角点
                for k in [1, 3]:  # 只遍历第2和第4个点
                    x, y, code_type, is_complete, weight = adjusted_labels[i + k][:5]
                    if is_complete:
                        update_heatmaps_and_direction(x, y, target_coords,
                                                      heatmap_codeX_24, heatmap_codeX_13, weight_codeX_24, weight,
                                                      heatmap_codeX_sin, heatmap_codeX_cos,
                                                      weight_codeX_sin, weight_codeX_cos, self.sigma)
            # 如果是QR码
            elif first_point_coords[2] == 1:
                # sigma = 2
                # 处理124角点
                for k in [0, 1, 3]:
                    # 读取当前角点的信息
                    x, y, code_type, is_complete, weight = adjusted_labels[i + k][:5]
                    if is_complete:
                        update_heatmaps_and_direction(x, y, target_coords,
                                                      heatmap_QR_124, heatmap_QR_3, weight_QR_124, weight,
                                                      heatmap_QR_sin, heatmap_QR_cos,
                                                      weight_QR_sin, weight_QR_cos, self.sigma)

                # 处理第三个角点
                x, y, code_type, is_complete, weight = adjusted_labels[i + 2][:5]
                if is_complete:
                    update_heatmaps_and_direction(x, y, target_coords,
                                                  heatmap_QR_3, heatmap_QR_124, weight_QR_3, weight,
                                                  heatmap_QR_sin, heatmap_QR_cos,
                                                  weight_QR_sin, weight_QR_cos, self.sigma)
            # 如果是DM码
            elif first_point_coords[2] == 2:
                # 处理134角点
                for k in [0, 2, 3]:
                    # 读取当前角点的信息
                    x, y, code_type, is_complete, weight = adjusted_labels[i + k][:5]
                    if is_complete:
                        update_heatmaps_and_direction(x, y, target_coords,
                                                      heatmap_DM_134, heatmap_DM_2, weight_DM_134, weight,
                                                      heatmap_DM_sin, heatmap_DM_cos,
                                                      weight_DM_sin, weight_DM_cos, self.sigma)

                # 处理第二个角点
                x, y, code_type, is_complete, weight = adjusted_labels[i + 1][:5]
                if is_complete:
                    update_heatmaps_and_direction(x, y, target_coords,
                                                  heatmap_DM_2, heatmap_DM_134, weight_DM_134, weight,
                                                  heatmap_DM_sin, heatmap_DM_cos,
                                                  weight_DM_sin, weight_DM_cos, self.sigma)
            # 如果是EAN
            elif first_point_coords[2] == 4:
                for k in [0, 1]:  # 只遍历第1和第2个点
                    x, y, code_type, is_complete, weight = adjusted_labels[i + k][:5]
                    if is_complete:
                        update_heatmaps_and_direction(x, y, target_coords,
                                                      heatmap_EAN_12, heatmap_EAN_34, weight_EAN_12, weight,
                                                      heatmap_EAN_sin, heatmap_EAN_cos,
                                                      weight_EAN_sin, weight_EAN_cos, self.sigma)

                # 处理24角点
                for k in [2, 3]:  # 只遍历第3和第4个点
                    x, y, code_type, is_complete, weight = adjusted_labels[i + k][:5]
                    if is_complete:
                        update_heatmaps_and_direction(x, y, target_coords,
                                                      heatmap_EAN_34, heatmap_EAN_12, weight_EAN_34, weight,
                                                      heatmap_EAN_sin, heatmap_EAN_cos,
                                                      weight_EAN_sin, weight_EAN_cos, self.sigma)


        heatmap = np.stack([
            heatmap_codeX_13, heatmap_codeX_24,  # Code128系列的位置热图
            heatmap_codeX_sin, heatmap_codeX_cos,  # Code128系列的方向热图
            heatmap_QR_124, heatmap_QR_3,  # QR的位置热图
            heatmap_QR_sin, heatmap_QR_cos,  # QR的方向热图
            heatmap_DM_134, heatmap_DM_2,  # DM的位置热图
            heatmap_DM_sin, heatmap_DM_cos,  # DM的方向热图
            heatmap_EAN_12, heatmap_EAN_34,  # EAN的位置热图
            heatmap_EAN_sin, heatmap_EAN_cos,  # EAN的方向热图
        ], axis=0)

        weight_map = np.stack([
            weight_codeX_13, weight_codeX_24,  # Code128的权重图
            weight_codeX_sin, weight_codeX_cos,  # Code128的方向权重图
            weight_QR_124, weight_QR_3,  # QR的权重图
            weight_QR_sin, weight_QR_cos,  # QR的方向权重图
            weight_DM_134, weight_DM_2,  # DM的权重图
            weight_DM_sin, weight_DM_cos,  # DM的方向权重图
            weight_EAN_12, weight_EAN_34,  # EAN的权重图
            weight_EAN_sin, weight_EAN_cos,  # EAN的方向权重图
        ], axis=0)

        return input_image, heatmap, adjusted_labels, weight_map, radius, image_path

def calculate_center_point(points):
    # 计算中心点，这里的points是四个角点的坐标列表
    sum_x, sum_y = 0, 0
    for point in points:
        sum_x += point[0]
        sum_y += point[1]
    center_x, center_y = sum_x / 4, sum_y / 4  # 四个角点的平均值
    return center_x, center_y

def resize_and_padding(image, target_size, fill_color=(114, 114, 114)):
    """
    将图片缩放并填充到目标尺寸。

    :param image: 原始PIL图像。
    :param target_size: 目标尺寸（宽度, 高度）。
    :param fill_color: 填充颜色，默认为灰色。

    :return: 缩放并填充后的图像，缩放比例，填充的像素数量。
    """
    # 检查图像模式，并为灰度图像调整填充颜色
    if image.mode == 'L':
        # 灰度图像使用单一颜色值填充
        if isinstance(fill_color, tuple):
            fill_color = fill_color[0]

    # 获取原始图像尺寸
    original_width, original_height = image.size
    # 获取放缩后的尺寸
    target_width, target_height = target_size

    # 计算缩放比例
    scale = min(target_width / original_width, target_height / original_height)

    # 计算缩放后的图像尺寸
    new_width = int(original_width * scale)
    new_height = int(original_height * scale)

    # 缩放图像
    image = image.resize((new_width, new_height), PIL.Image.Resampling.LANCZOS)

    # 计算填充的像素数量
    padding_left = (target_width - new_width) // 2
    padding_right = target_width - new_width - padding_left
    padding_top = (target_height - new_height) // 2
    padding_bottom = target_height - new_height - padding_top

    # 填充图像
    image = PIL.ImageOps.expand(image, (padding_left, padding_top, padding_right, padding_bottom), fill=fill_color)

    # 新图像的缩放比例和填充像素数量
    scale_x = new_width / original_width
    scale_y = new_height / original_height
    padding = (padding_left, padding_top, padding_right, padding_bottom)

    return image, (scale_x, scale_y), padding


def adjust_keypoints(labels, scale, padding):
    """
    根据图像的缩放和填充调整关键点的坐标。

    :param labels: 原始关键点标签列表，每个元素是一个 (x, y, classification, is_complete, weight) 五元组。
    :param scale: 缩放比例，是一个 (scale_x, scale_y) 元组。
    :param padding: 填充的像素数量，是一个 (left, top, right, bottom) 元组。

    :return: 调整后的关键点标签列表。
    """
    adjusted_labels = []
    scale_x, scale_y = scale
    padding_left, padding_top, _, _ = padding

    for label in labels:
        x, y, classification, is_complete, weight = label
        # 调整关键点坐标
        x = x * scale_x + padding_left
        y = y * scale_y + padding_top
        adjusted_label = (x, y, classification, is_complete, weight)  # 保持权重不变
        adjusted_labels.append(adjusted_label)

    return adjusted_labels

def generate_heatmap(heatmap, center, sigma, scale=1):
    """
    根据给定的中心点生成一个高斯热力图。

    :param heatmap: ndarray, 2D的numpy数组，用作初始化的热力图。
    :param center: tuple, 关键点的坐标，格式为 (x, y)。
    :param sigma: float, 高斯核的标准差，决定热力图的大小/扩散。
    :param scale: float, 可选参数，默认为2.0，用于调整sigma。

    :return: heatmap: ndarray, 更新后的热力图。radius: 高斯圈半径

    步骤:
    1. 获取热力图的大小。
    2. 调整σ值。
    3. 设置阈值和增量，用于确定高斯分布的边界。
    4. 遍历以关键点为中心的区域。
    5. 对于这个区域中的每个点，使用高斯公式更新其值。
    6. 返回更新后的热力图。
    """

    # 获取热力图尺寸并调整σ值
    h, w = heatmap.shape
    sigma *= scale
    # 设置经验阈值和增量
    th = 4.6052
    delta = math.sqrt(th * 2)
    # 高斯圈半径
    radius = sigma * delta
    # 获取关键点的整数坐标
    x0, y0 = int(center[0]), int(center[1])
    # 遍历关键点周围的区域
    for y in range(-int(radius), int(radius) + 1):
        for x in range(-int(radius), int(radius) + 1):
            # 确保不超出热力图边界
            if x0 + x >= 0 and x0 + x < w and y0 + y >= 0 and y0 + y < h:
                # 计算高斯函数的分子
                d = (x ** 2 + y ** 2) / (2 * sigma ** 2)
                exponent = d
                # 仅在阈值内更新热力图
                if exponent <= th:
                    heatmap[y0 + y, x0 + x] = max(heatmap[y0 + y, x0 + x], math.exp(-exponent))

    return heatmap,radius


def calculate_angle(p1, p2):
    # 计算两点间的角度，x轴正方向为0度
    dx, dy = p2[0] - p1[0], p2[1] - p1[1]
    angle_rad = np.arctan2(dy, dx) # 弧度
    return angle_rad


def collate_fn(batch):
    """
    将一个batch的数据转换为模型的输入。
    :param batch：一个batch的数据，是一个列表，每个元素是一个元组，包含了图像、关键点、权重、半径和图像路径。

    :return: 转换后的数据，包含了图像、关键点、权重、半径和图像路径。
    """
    images, heatmaps, labels, weight_maps, radii, image_paths = zip(*batch)

    images = [image.clone().detach() for image in images]
    heatmaps = [torch.tensor(heatmap) for heatmap in heatmaps]
    weight_maps = [torch.tensor(weight_map) for weight_map in weight_maps]
    radii = [torch.tensor(radius) for radius in radii]

    images = torch.stack(images)
    heatmaps = torch.stack(heatmaps)
    weight_maps = torch.stack(weight_maps)
    radii = torch.stack(radii)

    return images, heatmaps, labels, weight_maps, radii, image_paths


def visualize_heatmaps(image, heatmaps, score_threshold):
    """
    可视化缩放后的原图，并在其上绘制Code128和QR码的位置关键点和基于sin和cos热力图计算的角度。
    同时单独显示每个位置的热力图。

    :param image: 输入图像 (H, W, C)
    :param heatmaps: 关键点热力图 (8, H, W)，其中包含四个位置热力图和四个角度热力图
    :param score_threshold: 热力图中的阈值
    """
    num_rows = 6  # 1行原图 + 5行条码热图
    plt.figure(figsize=(16, 4 * num_rows))  # 每行高度设为4，总高度为行数乘以4

    # 缩放原图以匹配热力图的大小
    heatmap_height, heatmap_width = heatmaps[0].shape
    resized_image = cv2.resize(image, (heatmap_width, heatmap_height))

    # 展示原图
    ax_original = plt.subplot(num_rows, 4, (1, 4))
    ax_original.imshow(resized_image)
    ax_original.title.set_text("Original Image with Directions")
    ax_original.axis('off')

    # 在原图上绘制所有关键点的方向
    for i in range(2):  # 绘制Code128的方向
        y_coords, x_coords = np.where(heatmaps[i] >= score_threshold)
        print(y_coords, x_coords)
        for y, x in zip(y_coords, x_coords):
            sin_val = heatmaps[2][y, x]  # Code128的sin值
            cos_val = heatmaps[3][y, x]  # Code128的cos值
            angle_rad = math.atan2(sin_val, cos_val)
            # 计算箭头的终点坐标
            dx = x.item() + math.cos(angle_rad) * 30  # 箭头长度
            dy = y.item() + math.sin(angle_rad) * 30
            ax_original.arrow(x, y, dx - x, dy - y, head_width=3, head_length=3, fc='red', ec='red')

    for i in range(2):  # 绘制QR的方向
        y_coords, x_coords = np.where(heatmaps[i + 4] >= score_threshold)
        for y, x in zip(y_coords, x_coords):
            sin_val = heatmaps[6][y, x]  # QR的sin值
            cos_val = heatmaps[7][y, x]  # QR的cos值
            angle_rad = math.atan2(sin_val, cos_val)
            # 计算箭头的终点坐标
            dx = x.item() + math.cos(angle_rad) * 30  # 箭头长度
            dy = y.item() + math.sin(angle_rad) * 30
            ax_original.arrow(x, y, dx - x, dy - y, head_width=3, head_length=3, fc='red', ec='red')

    for i in range(2):  # 绘制DM的方向
        y_coords, x_coords = np.where(heatmaps[i + 8] >= score_threshold)
        for y, x in zip(y_coords, x_coords):
            sin_val = heatmaps[10][y, x]  # QR的sin值
            cos_val = heatmaps[11][y, x]  # QR的cos值
            angle_rad = math.atan2(sin_val, cos_val)
            # 计算箭头的终点坐标
            dx = x.item() + math.cos(angle_rad) * 30  # 箭头长度
            dy = y.item() + math.sin(angle_rad) * 30
            ax_original.arrow(x, y, dx - x, dy - y, head_width=3, head_length=3, fc='red', ec='red')

    for i in range(2):  # 绘制EAN的方向
        y_coords, x_coords = np.where(heatmaps[i + 12] >= score_threshold)
        for y, x in zip(y_coords, x_coords):
            sin_val = heatmaps[14][y, x]  # ean的sin值
            cos_val = heatmaps[15][y, x]  # EAN的cos值
            angle_rad = math.atan2(sin_val, cos_val)
            # 计算箭头的终点坐标
            dx = x.item() + math.cos(angle_rad) * 30  # 箭头长度
            dy = y.item() + math.sin(angle_rad) * 30
            ax_original.arrow(x, y, dx - x, dy - y, head_width=3, head_length=3, fc='red', ec='red')

    # 对于每种条码的处理，都需要调整索引以适应实际的行列数
    for barcode_index in range(4):  # 假设有5种条码
        for i in range(2):  # 每种条码两个热图
            ax = plt.subplot(num_rows, 4, barcode_index * 4 + 5 + i)
            pos_heatmap = heatmaps[barcode_index * 4 + i]
            ax.imshow(pos_heatmap, cmap='hot')
            barcode_name = ['CodeX', 'QR', 'DM', 'EAN'][barcode_index]
            ax.title.set_text(f"{barcode_name} Keypoint Heatmap {i + 1}")
            ax.axis('off')

    plt.tight_layout()
    plt.show()


# 用于测试的函数
def t_dataset():
    # 将图片转为(C*H*W)的tensor
    transform = transforms.ToTensor()
    dataset = KeypointDataset(r'../data/code_1/images/train',
                              r'../data/code_1/labels/train', transform=transform)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True, collate_fn=collate_fn)

    for batch_images, batch_heatmaps, batch_labels, batch_weight_map, batch_radius, image_path in dataloader:
        for image, heatmaps, labels, radius in zip(batch_images, batch_heatmaps, batch_labels,batch_radius):
            image_np = image.numpy().transpose(1, 2, 0)  # 将C, H, W格式转为H, W, C
            heatmaps_np = heatmaps.numpy()
            visualize_heatmaps(image_np, heatmaps_np, 1)
        break

if __name__ == "__main__":
    t_dataset()
